package com.example.vpnservicedemo.receive;

import android.net.VpnService;

import com.example.vpnservicedemo.ReadVpn;

import org.pcap4j.packet.IpPacket;
import org.pcap4j.packet.IpV4Packet;
import org.pcap4j.packet.IpV4Rfc1349Tos;
import org.pcap4j.packet.IpV6Packet;
import org.pcap4j.packet.Packet;
import org.pcap4j.packet.UdpPacket;
import org.pcap4j.packet.UnknownPacket;
import org.pcap4j.packet.namednumber.IpNumber;
import org.pcap4j.packet.namednumber.IpVersion;
import org.pcap4j.packet.namednumber.UdpPort;

import java.io.IOException;
import java.net.Inet4Address;
import java.net.Inet6Address;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.ProtocolFamily;
import java.net.StandardProtocolFamily;
import java.nio.ByteBuffer;
import java.nio.channels.DatagramChannel;

public class UdpPipe implements Pipe{

    private VpnService vpnService;

    private ReadVpn readVpn;

    private DatagramChannel channel;

    private IpPacket srcIppacket;


    public UdpPipe(VpnService vpnService, ReadVpn readVpn) {
        this.vpnService = vpnService;
        this.readVpn = readVpn;
    }

    @Override
    public int receive(IpPacket ipPacket) {

        try {

            ProtocolFamily protocolFamily = StandardProtocolFamily.INET;

            if (ipPacket instanceof IpV6Packet) {
                protocolFamily = StandardProtocolFamily.INET6;
            }

            UdpPacket udpPacket = (UdpPacket)ipPacket.getPayload();

            if(channel == null){
                srcIppacket = ipPacket;

                channel = DatagramChannel.open(protocolFamily);
                vpnService.protect(channel.socket());
                channel.configureBlocking(false);
                channel.connect(new InetSocketAddress(ipPacket.getHeader().getDstAddr(), udpPacket.getHeader().getDstPort().valueAsInt()));

                UdpSelector.registerRead(channel,this);
            }

            Packet payload = udpPacket.getPayload();
            if(payload!=null && payload.getRawData().length>0){

                if(channel.isConnected()){
                    channel.write(ByteBuffer.wrap(payload.getRawData()));
                }
            }


        } catch (Exception e) {
            System.out.println("UdpPipe receive 异常："+e.getMessage());
            close();
            return -1;
        }

        return 0;
    }

    public void close() {
        if(channel!=null){
            try {
                channel.close();
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }
    }

    public void pushData(byte[] data) {
        byte[] bytes = buildUdpPacket(data);
        readVpn.pushVpn(bytes);
    }

        private byte[] buildUdpPacket(byte[] data) {

            UdpPacket udpPacket = (UdpPacket) this.srcIppacket.getPayload();

        UdpPort srcPort = udpPacket.getHeader().getSrcPort();
        UdpPort dstPort = udpPacket.getHeader().getDstPort();


        UnknownPacket.Builder unknownBuilder = new UnknownPacket.Builder();
        unknownBuilder.rawData(data);

        //创建一个udp包
        UdpPacket.Builder udpBuilder = new UdpPacket.Builder()
                .srcPort(dstPort)
                .dstPort(srcPort)
                .payloadBuilder(unknownBuilder)
                .correctLengthAtBuild(true);


        InetAddress srcAddr = this.srcIppacket.getHeader().getSrcAddr();
        InetAddress dstAddr = this.srcIppacket.getHeader().getDstAddr();

        //创建一个ip包
        IpPacket newIpPacket;
        if (this.srcIppacket instanceof IpV6Packet) {
            newIpPacket = new IpV6Packet.Builder()
                    .version(IpVersion.IPV6)
                    .srcAddr((Inet6Address) dstAddr)
                    .dstAddr((Inet6Address) srcAddr)
                    .correctLengthAtBuild(true)
                    .payloadBuilder(udpBuilder)
                    .build();
        } else {

            newIpPacket = new IpV4Packet.Builder()
                    .version(IpVersion.IPV4)
                    .ttl((byte) 45)
                    .tos(IpV4Rfc1349Tos.newInstance((byte) 0))
                    .protocol(IpNumber.UDP)
                    .srcAddr((Inet4Address) dstAddr)
                    .dstAddr((Inet4Address) srcAddr)
                    .correctLengthAtBuild(true)
                    .correctChecksumAtBuild(true)
                    .payloadBuilder(udpBuilder)
                    .identification((byte) TcpPipe.getIpId())
                    .build();
        }
        return newIpPacket.getRawData();
    }
}
